/**
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/

/* !
 * \file whole_reduce_sum_custom.asc
 * \brief
 */

#include <cstdint>
#include <iostream>
#include <vector>
#include <algorithm>
#include <iterator>
#include "acl/acl.h"
#include "kernel_operator.h"
#include "data_utils.h"

constexpr uint32_t byteAlign = 32;

struct WholeReduceSumCustomTilingData {
    uint32_t totalLength;
    uint32_t rows;
    uint32_t cols;
};

__aicore__ inline uint32_t ceil_div(uint32_t a, uint32_t b)
{
    if (b == 0)
        return a;
    return (a + b - 1) / b;
}

template <typename datatype> class KernelWholeReduceSum {
public:
    __aicore__ inline KernelWholeReduceSum() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, WholeReduceSumCustomTilingData tilingData)
    {
        this->rows = tilingData.rows;
        this->cols = tilingData.cols;
        this->totalLength = tilingData.totalLength;
        this->colAligned = ceil_div(this->cols * sizeof(datatype), byteAlign) * byteAlign;
        this->totalSizeAligned = this->colAligned * this->rows;
        uint32_t reducedSizeAligned = ceil_div(this->rows * sizeof(datatype), byteAlign) * byteAlign;
        xGm.SetGlobalBuffer((__gm__ datatype *)x, this->totalLength);
        yGm.SetGlobalBuffer((__gm__ datatype *)y, this->rows);
        pipe.InitBuffer(inQueueX, 1, this->totalSizeAligned);
        pipe.InitBuffer(outQueueY, 1, reducedSizeAligned);
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<datatype> xLocal = inQueueX.AllocTensor<datatype>();
        uint32_t colBytes = this->cols * sizeof(datatype);
        AscendC::DataCopyExtParams copyParams = {(uint16_t)this->rows, colBytes, 0, 0, 0};
        uint8_t rpad = (this->colAligned - this->cols * sizeof(datatype)) / sizeof(datatype);
        AscendC::DataCopyPadExtParams<datatype> padParams = {false, 0, rpad, 0};
        AscendC::DataCopyPad<datatype>(xLocal, xGm, copyParams, padParams);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<datatype> xLocal = inQueueX.DeQue<datatype>();
        AscendC::LocalTensor<datatype> yLocal = outQueueY.AllocTensor<datatype>();

        int32_t mask = 256 / sizeof(datatype);
        int32_t srcStride = this->colAligned / byteAlign;
        AscendC::WholeReduceSum<datatype, true>(yLocal, xLocal, this->cols, this->rows, 1, 1, srcStride);
        outQueueY.EnQue<datatype>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<datatype> yLocal = outQueueY.DeQue<datatype>();
        uint16_t reducedBytes = this->rows * sizeof(datatype);
        AscendC::DataCopyExtParams copyParams = {1, reducedBytes, 0, 0, 0};
        AscendC::DataCopyPad<datatype>(yGm, yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outQueueY;
    AscendC::GlobalTensor<datatype> xGm;
    AscendC::GlobalTensor<datatype> yGm;
    uint32_t totalLength;
    uint32_t rows;
    uint32_t cols;
    uint32_t colAligned;
    uint32_t totalSizeAligned;
};

__aicore__ inline void CopyTiling(WholeReduceSumCustomTilingData *tiling, GM_ADDR tilingGM)
{
    uint32_t *ptr = reinterpret_cast<uint32_t *>(tiling);
    auto tiling32 = reinterpret_cast<__gm__ uint32_t *>(tilingGM);

    for (int i = 0; i < sizeof(WholeReduceSumCustomTilingData) / sizeof(uint32_t); i++, ptr++) {
        *ptr = *(tiling32 + i);
    }
    return;
}

__global__ __aicore__ void whole_reduce_sum_custom(GM_ADDR x, GM_ADDR y, GM_ADDR tiling)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    KernelWholeReduceSum<half> op;
    WholeReduceSumCustomTilingData tilingData;
    CopyTiling(&tilingData, tiling);
    op.Init(x, y, tilingData);
    op.Process();
}

constexpr uint32_t ROWS = 13;
constexpr uint32_t COLS = 123;
constexpr uint32_t TOTAL_SIZE = ROWS * COLS;
constexpr uint32_t TILING_SIZE = 12;

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 1;
    size_t typeSize = sizeof(uint16_t);
    size_t inputByteSize = ROWS * COLS * typeSize;
    size_t outputByteSize = ROWS * typeSize;
    size_t tilingSize = TILING_SIZE;

    aclInit(nullptr);
    int32_t deviceId = 0;
    aclrtSetDevice(deviceId);
    aclrtStream stream = nullptr;
    aclrtCreateStream(&stream);

    uint8_t *xHost, *yHost, *tilingHost;
    uint8_t *xDevice, *yDevice, *tilingDevice;

    aclrtMallocHost((void **)(&xHost), inputByteSize);
    aclrtMallocHost((void **)(&yHost), outputByteSize);
    aclrtMallocHost((void **)(&tilingHost), tilingSize);
    aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void **)&yDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void **)&tilingDevice, tilingSize, ACL_MEM_MALLOC_HUGE_FIRST);

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);
    ReadFile("./input/tiling.bin", tilingSize, tilingHost, tilingSize);

    aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(tilingDevice, tilingSize, tilingHost, tilingSize, ACL_MEMCPY_HOST_TO_DEVICE);

    whole_reduce_sum_custom<<<blockDim, nullptr, stream>>>(xDevice, yDevice, tilingDevice);
    aclrtSynchronizeStream(stream);

    aclrtMemcpy(yHost, outputByteSize, yDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST);
    WriteFile("./output/output.bin", yHost, outputByteSize);

    aclrtFree(xDevice);
    aclrtFree(yDevice);
    aclrtFree(tilingDevice);
    aclrtFreeHost(xHost);
    aclrtFreeHost(yHost);
    aclrtFreeHost(tilingHost);

    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();

    return 0;
}

